import sys
import os
import numpy as np
import cv2
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QFileDialog, QComboBox, QSlider, QGridLayout,
                             QSpinBox, QMessageBox)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QRect
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import pickle

# 导入自定义模块
from image_processing import ImageProcessor
from feature_extraction import FeatureExtractor
from embroidery_model import EmbroideryModel
from gui_components import GridViewer, FeatureVisualizer

class EmbroideryApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("刺绣针法分析系统")
        self.setGeometry(100, 100, 1200, 800)
        
        # 初始化变量
        self.image_path = None
        self.original_image = None
        self.grid_image = None
        self.grid_size = 30  # 默认网格大小
        self.current_stitch = "十字挑"  # 默认针法
        self.stitch_types = ["十字挑", "斜挑", "平挑", "横挑", "竖挑"]
        self.model = None
        self.features = None
        self.grid_features = None
        
        # 初始化处理器
        self.image_processor = ImageProcessor()
        self.feature_extractor = FeatureExtractor()
        self.embroidery_model = EmbroideryModel()
        
        # 加载模型（如果存在）
        self.load_model()
        
        # 设置UI
        self.init_ui()
    
    def init_ui(self):
        # 主布局
        central_widget = QWidget()
        main_layout = QHBoxLayout()
        
        # 左侧控制面板
        control_panel = QWidget()
        control_layout = QVBoxLayout()
        
        # 图片加载按钮
        self.load_button = QPushButton("加载图片")
        self.load_button.clicked.connect(self.load_image)
        control_layout.addWidget(self.load_button)
        
        # 网格大小控制
        grid_control = QWidget()
        grid_layout = QHBoxLayout()
        grid_layout.addWidget(QLabel("网格大小:"))
        self.grid_size_spin = QSpinBox()
        self.grid_size_spin.setRange(10, 100)
        self.grid_size_spin.setValue(self.grid_size)
        self.grid_size_spin.valueChanged.connect(self.update_grid_size)
        grid_layout.addWidget(self.grid_size_spin)
        grid_control.setLayout(grid_layout)
        control_layout.addWidget(grid_control)
        
        # 针法选择
        stitch_control = QWidget()
        stitch_layout = QHBoxLayout()
        stitch_layout.addWidget(QLabel("针法选择:"))
        self.stitch_combo = QComboBox()
        self.stitch_combo.addItems(self.stitch_types)
        self.stitch_combo.currentTextChanged.connect(self.update_stitch)
        stitch_layout.addWidget(self.stitch_combo)
        stitch_control.setLayout(stitch_layout)
        control_layout.addWidget(stitch_control)
        
        # 特征提取按钮
        self.extract_button = QPushButton("提取特征")
        self.extract_button.clicked.connect(self.extract_features)
        control_layout.addWidget(self.extract_button)
        
        # 生成刺绣图按钮
        self.generate_button = QPushButton("生成刺绣图")
        self.generate_button.clicked.connect(self.generate_embroidery)
        control_layout.addWidget(self.generate_button)
        
        # 训练模型按钮
        self.train_button = QPushButton("训练模型")
        self.train_button.clicked.connect(self.train_model)
        control_layout.addWidget(self.train_button)
        
        # 保存结果按钮
        self.save_button = QPushButton("保存结果")
        self.save_button.clicked.connect(self.save_result)
        control_layout.addWidget(self.save_button)
        
        control_panel.setLayout(control_layout)
        control_panel.setFixedWidth(200)
        
        # 右侧显示区域
        display_panel = QWidget()
        display_layout = QVBoxLayout()
        
        # 原图和网格图显示
        image_display = QWidget()
        image_layout = QHBoxLayout()
        
        # 原图显示
        self.original_label = QLabel("原始图像")
        self.original_label.setAlignment(Qt.AlignCenter)
        self.original_label.setMinimumSize(400, 300)
        self.original_label.setStyleSheet("border: 1px solid black")
        image_layout.addWidget(self.original_label)
        
        # 网格图显示
        self.grid_label = QLabel("网格图像")
        self.grid_label.setAlignment(Qt.AlignCenter)
        self.grid_label.setMinimumSize(400, 300)
        self.grid_label.setStyleSheet("border: 1px solid black")
        image_layout.addWidget(self.grid_label)
        
        image_display.setLayout(image_layout)
        display_layout.addWidget(image_display)
        
        # 特征可视化和结果显示
        result_display = QWidget()
        result_layout = QHBoxLayout()
        
        # 特征可视化
        self.feature_label = QLabel("特征可视化")
        self.feature_label.setAlignment(Qt.AlignCenter)
        self.feature_label.setMinimumSize(400, 300)
        self.feature_label.setStyleSheet("border: 1px solid black")
        result_layout.addWidget(self.feature_label)
        
        # 刺绣结果显示
        self.result_label = QLabel("刺绣结果")
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setMinimumSize(400, 300)
        self.result_label.setStyleSheet("border: 1px solid black")
        result_layout.addWidget(self.result_label)
        
        result_display.setLayout(result_layout)
        display_layout.addWidget(result_display)
        
        display_panel.setLayout(display_layout)
        
        # 添加到主布局
        main_layout.addWidget(control_panel)
        main_layout.addWidget(display_panel)
        
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)
    
    def load_image(self):
        file_dialog = QFileDialog()
        image_path, _ = file_dialog.getOpenFileName(self, "选择图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp)")
        
        if image_path:
            self.image_path = image_path
            self.original_image = cv2.imread(image_path)
            self.original_image = cv2.cvtColor(self.original_image, cv2.COLOR_BGR2RGB)
            
            # 显示原图
            self.display_image(self.original_image, self.original_label)
            
            # 生成并显示网格图
            self.update_grid()
    
    def update_grid_size(self, value):
        self.grid_size = value
        if self.original_image is not None:
            self.update_grid()
    
    def update_grid(self):
        if self.original_image is not None:
            self.grid_image = self.image_processor.create_grid(self.original_image, self.grid_size)
            self.display_image(self.grid_image, self.grid_label)
    
    def update_stitch(self, stitch_type):
        self.current_stitch = stitch_type
    
    def extract_features(self):
        if self.original_image is not None:
            # 提取特征
            self.grid_features = self.feature_extractor.extract_grid_features(
                self.original_image, self.grid_size)
            
            # 可视化特征
            feature_viz = self.feature_extractor.visualize_features(self.original_image, 
                                                                  self.grid_features, 
                                                                  self.grid_size)
            self.display_image(feature_viz, self.feature_label)
            
            QMessageBox.information(self, "特征提取", "特征提取完成！")
    
    def generate_embroidery(self):
        if self.original_image is not None and self.grid_features is not None:
            # 使用模型预测或直接应用针法
            if self.model is not None:
                # 使用训练好的模型预测
                embroidery_image = self.embroidery_model.apply_model_prediction(
                    self.original_image, self.grid_features, self.model, self.grid_size)
            else:
                # 直接应用选定的针法
                embroidery_image = self.embroidery_model.apply_stitch(
                    self.original_image, self.current_stitch, self.grid_size)
            
            self.display_image(embroidery_image, self.result_label)
    
    def train_model(self):
        # 这里应该打开一个对话框，让用户选择训练数据目录
        # 简化起见，我们假设用户已经准备好了训练数据
        QMessageBox.information(self, "模型训练", "请准备好训练数据，包含图像和对应的针法标注。")
        
        train_dir = QFileDialog.getExistingDirectory(self, "选择训练数据目录")
        if train_dir:
            try:
                # 训练模型
                self.model = self.embroidery_model.train_model(train_dir)
                # 保存模型
                self.save_model()
                QMessageBox.information(self, "模型训练", "模型训练完成并保存！")
            except Exception as e:
                QMessageBox.warning(self, "训练错误", f"训练过程中出错: {str(e)}")
    
    def save_model(self):
        if self.model is not None:
            model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", "embroidery_model.pkl")
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            with open(model_path, 'wb') as f:
                pickle.dump(self.model, f)
    
    def load_model(self):
        model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", "embroidery_model.pkl")
        if os.path.exists(model_path):
            try:
                with open(model_path, 'rb') as f:
                    self.model = pickle.load(f)
            except Exception as e:
                print(f"加载模型出错: {str(e)}")
    
    def save_result(self):
        if hasattr(self, 'result_image'):
            file_dialog = QFileDialog()
            save_path, _ = file_dialog.getSaveFileName(self, "保存结果", "", "图片文件 (*.png *.jpg)")
            if save_path:
                cv2.imwrite(save_path, cv2.cvtColor(self.result_image, cv2.COLOR_RGB2BGR))
                QMessageBox.information(self, "保存结果", "结果已保存！")
    
    def display_image(self, image, label):
        h, w, c = image.shape
        bytes_per_line = 3 * w
        q_image = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_image)
        
        # 调整大小以适应标签
        label_size = label.size()
        pixmap = pixmap.scaled(label_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        
        label.setPixmap(pixmap)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = EmbroideryApp()
    window.show()
    sys.exit(app.exec_())